# -*- coding: UTF-8 -*-

from llm.llm import llm

# from tools.invoke_api import *
# from trace_agent.trace_agent import *
# from utils.utils import *


from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)


summary = """
"header": {
    "title": "Ubuntu资产遭遇高危反弹Shell攻击",
    "desc": "taomingwei-node3-ubuntu自2022年4月3日至2025年4月8日共遭受13次网络攻击，包括检测到Linux可疑命令和高风险的perl反弹shell事件。建议立即排查并采取措施防止进一步恶意活动。"
}
"""


def main():
    sys_tmpl = """
    根据内容创建一个美观的Vue 3 Setup语法组件：

    布局结构：
    使用全局<app-container>组件作为最外层包裹，不需要import

    要求：
    使用现代UI设计原则
    一致的配色方案
    清晰的视觉层次
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    # print("chat_prompt", chat_prompt.content)

    chain = chat_prompt | llm

    res = chain.invoke({"summary": summary})

    print("res", res)

    with open("test.html", "w", encoding="utf-8") as f:
        f.write(res.text())

    # 保存结果

    print(res.text())


main()
